from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

# 中文文档 http://backtrader.com.cn/docu/#1
# 英文文档 https://www.backtrader.com/docu/

import datetime  # 日期时间类
import os.path  # 路径管理
import sys  # 获取当前运行脚本的路径

# 导入backtrader框架
import backtrader as bt

import pandas as pd

# 创建自己的策略
class MyStrategy(bt.Strategy):
    # 添加参数 自定义策略
    # 调用时 cerebro.addstrategy(MyStrategy, exitbars=7)
    # 寻找最优参数 cerebro.optstrategy(MyStrategy, exitbars=range(3,15))
    params = (
        # 持仓够5个单位就卖出
        ('exitbars', 5),
    )


    # 定义日志格式
    def log(self, txt, dt=None):
        # 记录策略的执行日志
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    # 初始化相关数据
    def __init__(self):
        # 保存收盘价的引用
        self.dataclose = self.datas[0].close

        # 当前订单状态 跟踪挂单
        self.order = None

        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None

    def notify_order(self, order):
        self.log("notify_order")
        if order.status in [order.Submitted, order.Accepted]:
            # broker 提交/接受 买/卖订单则什么都不做
            return

        # 检查一个订单是否完成  注意 当资金不足时 broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log('已买入, %.2f, 费用: %.2f, 佣金 %.2f' % (order.executed.price, order.executed.value, order.executed.comm))
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm 
            elif order.issell():
                self.log('已卖出, %.2f, 费用: %.2f, 佣金 %.2f' % (order.executed.price, order.executed.value, order.executed.comm))

            # 记录当前交易数量
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消/保证金不足/拒绝')

        # 其他状态记录为 无挂起订单
        self.order = None

    # 交易状态通知 一买一卖计算交易
    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log("交易利润, 毛利润 %.2f, 净利润 %.2f" % (trade.pnl, trade.pnlcomm))

    # 打印出行情数据
    def next(self):
        # 记录收盘价
        self.log('Close, %.2f' % self.dataclose[0])

        # 如果有订单正在挂起 不操作
        if self.order:
            return

        # 如果没有持仓则买入
        if not self.position:
            # 添加买入逻辑 如果k线收盘出现三连跌 则买入
            # 今天的收盘价 < 昨天收盘价
            if self.dataclose[0] < self.dataclose[-1]:
                    # 昨天收盘价 < 前天的收盘价
                    if self.dataclose[-1] < self.dataclose[-2]:
                        # 买入
                        self.log('买入, %.2f' % self.dataclose[0])
                        # 跟踪订单避免重复
                        self.order = self.buy()
        else:
            # 如果已经持仓 且当前交易数据量在买入后5个单位后
            # 将5替换为参数
            if len(self) >= (self.bar_executed + self.params.exitbars):
                # 全部卖出
                self.log('创建卖出, %.2f' % self.dataclose[0])
                # 跟踪订单避免重复
                self.order = self.sell()

if __name__ == '__main__':
    # 创建cerebro引擎
    cerebro = bt.Cerebro()
    cerebro.broker.set_cash(100000.00)  # 设置初始资金金额
    cerebro = bt.Cerebro(stdstats=False)
    # 账户信息
    cerebro.addobserver(bt.observers.Broker)
    # 买卖点
    cerebro.addobserver(bt.observers.BuySell)
    # 计算回撤
    cerebro.addobserver(bt.observers.DrawDown)
    # 资金曲线
    cerebro.addobserver(bt.observers.Value)
    # 期间收益
    cerebro.addobserver(bt.observers.TimeReturn)

    # cerebro引擎在后台创建broker(经纪人) 系统默认资金量为10000
    init_fund = cerebro.broker.getvalue()
    print(f'初始资金:{init_fund}')
    filename = 'sz000002.csv'
    print(filename)
    data = pd.read_csv(filename, index_col="date", parse_dates=True)
    load_data = bt.feeds.PandasData(dataname=data, fromdate=datetime.datetime(2020, 1, 1), todate=datetime.datetime(2020, 5, 1))
    # 加载交易数据
    cerebro.adddata(load_data)
    # 添加策略
    cerebro.addstrategy(MyStrategy)

    # 设置佣金0.001 除以100去掉%号
    cerebro.broker.setcommission(commission=0.001)

    cerebro.run()
    end_fund = cerebro.broker.getvalue()
    print(f'期末资金:{end_fund}')
    cerebro.plot(style="candle")